"""
事务管理器
负责事务的创建、提交、回滚和并发控制
"""

import threading
import time
import uuid
from typing import Dict, List, Optional, Set, Any
from enum import Enum
from dataclasses import dataclass
from collections import defaultdict

from .types import IsolationLevel

class TransactionState(Enum):
    """事务状态"""
    ACTIVE = "active"
    COMMITTED = "committed"
    ABORTED = "aborted"
    PREPARING = "preparing"
    PREPARED = "prepared"

# IsolationLevel 已移动到 types.py

class LockType(Enum):
    """锁类型"""
    SHARED = "shared"
    EXCLUSIVE = "exclusive"
    INTENT_SHARED = "intent_shared"
    INTENT_EXCLUSIVE = "intent_exclusive"

@dataclass
class Lock:
    """锁对象"""
    resource_id: str
    lock_type: LockType
    transaction_id: str
    acquired_time: float
    
    def is_compatible(self, other_lock: 'Lock') -> bool:
        """检查锁兼容性"""
        if self.transaction_id == other_lock.transaction_id:
            return True
        
        # 锁兼容性矩阵
        compatibility_matrix = {
            (LockType.SHARED, LockType.SHARED): True,
            (LockType.SHARED, LockType.EXCLUSIVE): False,
            (LockType.SHARED, LockType.INTENT_SHARED): True,
            (LockType.SHARED, LockType.INTENT_EXCLUSIVE): False,
            (LockType.EXCLUSIVE, LockType.SHARED): False,
            (LockType.EXCLUSIVE, LockType.EXCLUSIVE): False,
            (LockType.EXCLUSIVE, LockType.INTENT_SHARED): False,
            (LockType.EXCLUSIVE, LockType.INTENT_EXCLUSIVE): False,
            (LockType.INTENT_SHARED, LockType.SHARED): True,
            (LockType.INTENT_SHARED, LockType.EXCLUSIVE): False,
            (LockType.INTENT_SHARED, LockType.INTENT_SHARED): True,
            (LockType.INTENT_SHARED, LockType.INTENT_EXCLUSIVE): True,
            (LockType.INTENT_EXCLUSIVE, LockType.SHARED): False,
            (LockType.INTENT_EXCLUSIVE, LockType.EXCLUSIVE): False,
            (LockType.INTENT_EXCLUSIVE, LockType.INTENT_SHARED): True,
            (LockType.INTENT_EXCLUSIVE, LockType.INTENT_EXCLUSIVE): True,
        }
        
        return compatibility_matrix.get((self.lock_type, other_lock.lock_type), False)

@dataclass
class Transaction:
    """事务对象"""
    transaction_id: str
    start_time: float
    isolation_level: IsolationLevel
    state: TransactionState = TransactionState.ACTIVE
    locks: Set[str] = None
    read_set: Set[str] = None
    write_set: Set[str] = None
    undo_log: List[Dict[str, Any]] = None
    
    def __post_init__(self):
        if self.locks is None:
            self.locks = set()
        if self.read_set is None:
            self.read_set = set()
        if self.write_set is None:
            self.write_set = set()
        if self.undo_log is None:
            self.undo_log = []

class DeadlockDetector:
    """死锁检测器"""
    
    def __init__(self):
        self.wait_for_graph: Dict[str, Set[str]] = defaultdict(set)
        self.lock = threading.RLock()
    
    def add_wait_edge(self, waiting_tx: str, holding_tx: str):
        """添加等待边"""
        with self.lock:
            self.wait_for_graph[waiting_tx].add(holding_tx)
    
    def remove_wait_edge(self, waiting_tx: str, holding_tx: str):
        """移除等待边"""
        with self.lock:
            if waiting_tx in self.wait_for_graph:
                self.wait_for_graph[waiting_tx].discard(holding_tx)
    
    def detect_deadlock(self) -> Optional[List[str]]:
        """检测死锁，返回死锁环"""
        with self.lock:
            visited = set()
            rec_stack = set()
            
            for node in self.wait_for_graph:
                if node not in visited:
                    cycle = self._dfs_detect_cycle(node, visited, rec_stack, [])
                    if cycle:
                        return cycle
            
            return None
    
    def _dfs_detect_cycle(self, node: str, visited: Set[str], 
                         rec_stack: Set[str], path: List[str]) -> Optional[List[str]]:
        """DFS检测环"""
        visited.add(node)
        rec_stack.add(node)
        path.append(node)
        
        for neighbor in self.wait_for_graph[node]:
            if neighbor not in visited:
                cycle = self._dfs_detect_cycle(neighbor, visited, rec_stack, path)
                if cycle:
                    return cycle
            elif neighbor in rec_stack:
                # 找到环
                cycle_start = path.index(neighbor)
                return path[cycle_start:] + [neighbor]
        
        rec_stack.remove(node)
        path.pop()
        return None
    
    def clear_transaction(self, transaction_id: str):
        """清理事务相关的等待图"""
        with self.lock:
            # 移除该事务的所有等待边
            if transaction_id in self.wait_for_graph:
                del self.wait_for_graph[transaction_id]
            
            # 移除其他事务对该事务的等待
            for tx_id in self.wait_for_graph:
                self.wait_for_graph[tx_id].discard(transaction_id)

class LockManager:
    """锁管理器"""
    
    def __init__(self):
        self.lock_table: Dict[str, List[Lock]] = defaultdict(list)
        self.waiting_queue: Dict[str, List[Lock]] = defaultdict(list)
        self.transaction_locks: Dict[str, Set[str]] = defaultdict(set)
        self.lock = threading.RLock()
        self.deadlock_detector = DeadlockDetector()
    
    def acquire_lock(self, resource_id: str, lock_type: LockType, 
                    transaction_id: str, timeout: float = 30.0) -> bool:
        """获取锁"""
        with self.lock:
            new_lock = Lock(resource_id, lock_type, transaction_id, time.time())
            
            # 检查是否可以立即获得锁
            if self._can_acquire_lock(new_lock):
                self.lock_table[resource_id].append(new_lock)
                self.transaction_locks[transaction_id].add(resource_id)
                return True
            
            # 不能立即获得锁，加入等待队列
            self.waiting_queue[resource_id].append(new_lock)
            
            # 添加到等待图
            for existing_lock in self.lock_table[resource_id]:
                if not new_lock.is_compatible(existing_lock):
                    self.deadlock_detector.add_wait_edge(transaction_id, existing_lock.transaction_id)
            
            # 检测死锁
            deadlock_cycle = self.deadlock_detector.detect_deadlock()
            if deadlock_cycle and transaction_id in deadlock_cycle:
                # 发现死锁，移除等待请求
                self.waiting_queue[resource_id].remove(new_lock)
                return False
        
        # 等待锁被释放
        return self._wait_for_lock(new_lock, timeout)
    
    def _can_acquire_lock(self, new_lock: Lock) -> bool:
        """检查是否可以获得锁"""
        for existing_lock in self.lock_table[new_lock.resource_id]:
            if not new_lock.is_compatible(existing_lock):
                return False
        return True
    
    def _wait_for_lock(self, lock: Lock, timeout: float) -> bool:
        """等待锁"""
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            with self.lock:
                if self._can_acquire_lock(lock):
                    # 从等待队列移除并添加到锁表
                    if lock in self.waiting_queue[lock.resource_id]:
                        self.waiting_queue[lock.resource_id].remove(lock)
                    
                    self.lock_table[lock.resource_id].append(lock)
                    self.transaction_locks[lock.transaction_id].add(lock.resource_id)
                    
                    # 从等待图中移除
                    for existing_lock in self.lock_table[lock.resource_id]:
                        if existing_lock != lock:
                            self.deadlock_detector.remove_wait_edge(
                                lock.transaction_id, existing_lock.transaction_id)
                    
                    return True
            
            time.sleep(0.01)  # 短暂等待
        
        # 超时，移除等待请求
        with self.lock:
            if lock in self.waiting_queue[lock.resource_id]:
                self.waiting_queue[lock.resource_id].remove(lock)
        
        return False
    
    def release_lock(self, resource_id: str, transaction_id: str):
        """释放锁"""
        with self.lock:
            # 移除事务的锁
            locks_to_remove = []
            for lock in self.lock_table[resource_id]:
                if lock.transaction_id == transaction_id:
                    locks_to_remove.append(lock)
            
            for lock in locks_to_remove:
                self.lock_table[resource_id].remove(lock)
            
            self.transaction_locks[transaction_id].discard(resource_id)
            
            # 检查等待队列中是否有可以获得锁的请求
            self._process_waiting_queue(resource_id)
    
    def release_all_locks(self, transaction_id: str):
        """释放事务的所有锁"""
        with self.lock:
            resource_ids = self.transaction_locks[transaction_id].copy()
            
            for resource_id in resource_ids:
                self.release_lock(resource_id, transaction_id)
            
            # 清理死锁检测器
            self.deadlock_detector.clear_transaction(transaction_id)
    
    def _process_waiting_queue(self, resource_id: str):
        """处理等待队列"""
        waiting_locks = self.waiting_queue[resource_id].copy()
        
        for waiting_lock in waiting_locks:
            if self._can_acquire_lock(waiting_lock):
                self.waiting_queue[resource_id].remove(waiting_lock)
                self.lock_table[resource_id].append(waiting_lock)
                self.transaction_locks[waiting_lock.transaction_id].add(resource_id)

class TransactionManager:
    """事务管理器"""
    
    def __init__(self):
        self.active_transactions: Dict[str, Transaction] = {}
        self.lock_manager = LockManager()
        self.next_transaction_id = 1
        self.lock = threading.RLock()
        
        # 事务日志
        self.transaction_log: List[Dict[str, Any]] = []
        
        # 检查点管理
        self.last_checkpoint_time = time.time()
        self.checkpoint_interval = 300  # 5分钟
    
    def begin_transaction(self, isolation_level: IsolationLevel = IsolationLevel.READ_COMMITTED) -> str:
        """开始事务"""
        with self.lock:
            transaction_id = f"tx_{self.next_transaction_id}"
            self.next_transaction_id += 1
            
            transaction = Transaction(
                transaction_id=transaction_id,
                start_time=time.time(),
                isolation_level=isolation_level
            )
            
            self.active_transactions[transaction_id] = transaction
            
            # 记录日志
            self._log_transaction_event(transaction_id, "BEGIN", {
                "isolation_level": isolation_level.value
            })
            
            return transaction_id
    
    def commit_transaction(self, transaction_id: str) -> bool:
        """提交事务"""
        with self.lock:
            if transaction_id not in self.active_transactions:
                return False
            
            transaction = self.active_transactions[transaction_id]
            
            if transaction.state != TransactionState.ACTIVE:
                return False
            
            try:
                # 设置状态为准备提交
                transaction.state = TransactionState.PREPARING
                
                # 执行提交前检查
                if not self._pre_commit_validation(transaction):
                    self.abort_transaction(transaction_id)
                    return False
                
                # 写入提交日志
                self._log_transaction_event(transaction_id, "COMMIT", {
                    "read_set": list(transaction.read_set),
                    "write_set": list(transaction.write_set)
                })
                
                # 设置状态为已提交
                transaction.state = TransactionState.COMMITTED
                
                # 释放所有锁
                self.lock_manager.release_all_locks(transaction_id)
                
                # 清理事务
                del self.active_transactions[transaction_id]
                
                return True
                
            except Exception as e:
                # 提交失败，回滚事务
                self.abort_transaction(transaction_id)
                return False
    
    def abort_transaction(self, transaction_id: str) -> bool:
        """回滚事务"""
        with self.lock:
            if transaction_id not in self.active_transactions:
                return False
            
            transaction = self.active_transactions[transaction_id]
            
            try:
                # 设置状态为中止
                transaction.state = TransactionState.ABORTED
                
                # 执行回滚操作
                self._rollback_changes(transaction)
                
                # 记录回滚日志
                self._log_transaction_event(transaction_id, "ABORT", {
                    "reason": "explicit_abort"
                })
                
                # 释放所有锁
                self.lock_manager.release_all_locks(transaction_id)
                
                # 清理事务
                del self.active_transactions[transaction_id]
                
                return True
                
            except Exception as e:
                # 回滚失败，记录错误
                self._log_transaction_event(transaction_id, "ABORT_FAILED", {
                    "error": str(e)
                })
                return False
    
    def _pre_commit_validation(self, transaction: Transaction) -> bool:
        """提交前验证"""
        # 对于可序列化隔离级别，需要验证读集合
        if transaction.isolation_level == IsolationLevel.SERIALIZABLE:
            return self._validate_serializable(transaction)
        
        return True
    
    def _validate_serializable(self, transaction: Transaction) -> bool:
        """验证可序列化性"""
        # 简化的验证：检查读集合是否被其他事务修改
        for resource_id in transaction.read_set:
            for other_tx_id, other_tx in self.active_transactions.items():
                if other_tx_id != transaction.transaction_id:
                    if resource_id in other_tx.write_set:
                        if other_tx.start_time > transaction.start_time:
                            return False
        
        return True
    
    def _rollback_changes(self, transaction: Transaction):
        """回滚事务的修改"""
        # 按逆序执行撤销操作
        for undo_entry in reversed(transaction.undo_log):
            try:
                self._apply_undo_entry(undo_entry)
            except Exception as e:
                # 记录撤销失败，但继续处理其他撤销项
                self._log_transaction_event(transaction.transaction_id, "UNDO_FAILED", {
                    "undo_entry": undo_entry,
                    "error": str(e)
                })
    
    def _apply_undo_entry(self, undo_entry: Dict[str, Any]):
        """应用撤销项"""
        # 这里应该根据undo_entry的类型执行相应的撤销操作
        # 简化实现
        pass
    
    def _log_transaction_event(self, transaction_id: str, event_type: str, data: Dict[str, Any]):
        """记录事务事件"""
        log_entry = {
            "timestamp": time.time(),
            "transaction_id": transaction_id,
            "event_type": event_type,
            "data": data
        }
        
        self.transaction_log.append(log_entry)
        
        # 检查是否需要创建检查点
        if time.time() - self.last_checkpoint_time > self.checkpoint_interval:
            self._create_checkpoint()
    
    def _create_checkpoint(self):
        """创建检查点"""
        with self.lock:
            checkpoint_data = {
                "timestamp": time.time(),
                "active_transactions": list(self.active_transactions.keys()),
                "log_size": len(self.transaction_log)
            }
            
            # 简化实现：只记录检查点信息
            self._log_transaction_event("SYSTEM", "CHECKPOINT", checkpoint_data)
            
            self.last_checkpoint_time = time.time()
    
    def get_transaction(self, transaction_id: str) -> Optional[Transaction]:
        """获取事务对象"""
        return self.active_transactions.get(transaction_id)
    
    def get_active_transactions(self) -> List[str]:
        """获取活动事务列表"""
        with self.lock:
            return list(self.active_transactions.keys())
    
    def get_transaction_stats(self) -> Dict[str, Any]:
        """获取事务统计信息"""
        with self.lock:
            stats = {
                "active_transactions": len(self.active_transactions),
                "total_log_entries": len(self.transaction_log),
                "last_checkpoint": self.last_checkpoint_time,
                "deadlock_detection_enabled": True
            }
            
            # 按隔离级别统计
            isolation_stats = defaultdict(int)
            for tx in self.active_transactions.values():
                isolation_stats[tx.isolation_level.value] += 1
            
            stats["isolation_level_distribution"] = dict(isolation_stats)
            
            return stats
    
    def cleanup_old_logs(self, retention_hours: int = 24):
        """清理旧日志"""
        cutoff_time = time.time() - (retention_hours * 3600)
        
        with self.lock:
            self.transaction_log = [
                log for log in self.transaction_log 
                if log["timestamp"] > cutoff_time
            ]

# 全局事务管理器实例
_global_transaction_manager = None

def get_transaction_manager() -> TransactionManager:
    """获取全局事务管理器实例"""
    global _global_transaction_manager
    if _global_transaction_manager is None:
        _global_transaction_manager = TransactionManager()
    return _global_transaction_manager

def set_transaction_manager(manager: TransactionManager):
    """设置全局事务管理器实例"""
    global _global_transaction_manager
    _global_transaction_manager = manager


